#include "Proxy.h"
#include "debug.h"

#define WORKER_THREADS_PER_PROCESSOR 2
#define WAIT_TIMEOUT_INTERVAL   100     // ???
static const u_short _port = 5005;


Proxy::Proxy()
    : _iocp(NULL), _shutdownEvent(NULL), _listenSocket(INVALID_SOCKET)
{
}

Proxy::~Proxy()
{
    stop();
}

bool Proxy::run()
{
    if( _initialize() ) 
    {
        if( _createListenSocket() ) 
        {
            SYSTEM_INFO si = {0};
            GetSystemInfo(&si);
            _threads.resize(si.dwNumberOfProcessors * WORKER_THREADS_PER_PROCESSOR);

            DWORD threadID;
            for(size_t i = 0; i < _threads.size(); ++i)
            {
                _threads[i] = CreateThread(0, 0, &Proxy::_workerThreadProc, (void*)this, 0, &threadID);
            }

            return true;
        }
    }

    _uninitialize();
    return false;
}

void Proxy::stop()
{
    SetEvent(_shutdownEvent);

    for(size_t i = 0; i < _threads.size(); ++i)
    {
        PostQueuedCompletionStatus(_iocp, 0, (DWORD) NULL, NULL);
    }
    WaitForMultipleObjects(_threads.size(), &_threads[0], TRUE, INFINITE);

    _uninitialize();
}

bool Proxy::_initialize()
{
    WSADATA wsaData = {0};
    int res = WSAStartup(MAKEWORD(2,2), &wsaData);
    if(NO_ERROR != res) {
        dbgprint("Error occurred while executing WSAStartup(): %d", res);
        return false;
    }

    _shutdownEvent = CreateEvent(NULL, TRUE, FALSE, NULL);

    _iocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
    if(_iocp == NULL) {
        dbgprint("Error occurred while creating IOCP: %d.", WSAGetLastError());
        return false;
    }

    return true;
}

void Proxy::_uninitialize()
{
    if(_iocp != NULL) {
        CloseHandle(_shutdownEvent);
        _iocp = NULL;
    }

    if(_shutdownEvent != NULL) {
        CloseHandle(_shutdownEvent);
        _shutdownEvent = NULL;
    }

    if(_listenSocket != INVALID_SOCKET) {
        closesocket(_listenSocket);
        _listenSocket = INVALID_SOCKET;
    }

    _clearConnections();

    WSACleanup();
}

bool Proxy::_createListenSocket()
{
    _listenSocket = WSASocket(AF_INET, SOCK_STREAM, 0/*IPPROTO_IP*/, NULL, 0, WSA_FLAG_OVERLAPPED);
    if(_listenSocket == INVALID_SOCKET) {
        dbgprint("Error occurred while opening socket: %d.", WSAGetLastError());
        return false;
    }

    sockaddr_in serverAddress = {0};
    serverAddress.sin_family = AF_INET;
    serverAddress.sin_addr.s_addr = INADDR_ANY;
    serverAddress.sin_port = htons(_port);
    int res = bind(_listenSocket, (struct sockaddr*)&serverAddress, sizeof(serverAddress));
    if(res == SOCKET_ERROR) {
        dbgprint("Error occurred while binding: %d.", WSAGetLastError());
        return false;
    }
    
    res = listen(_listenSocket, SOMAXCONN);
    if(res == SOCKET_ERROR) {
        dbgprint("Error occurred while listening: %d", WSAGetLastError());
        return false;
    }

    HANDLE temp = CreateIoCompletionPort((HANDLE)_listenSocket, _iocp, 0, 0);
    if(temp == NULL) {
        dbgprint("Error occurred while executing CreateIoCompletionPort(): %d", GetLastError());
        return false;
    }

    return (_createConnection() != NULL);
}

DWORD WINAPI Proxy::_workerThreadProc(LPVOID lpParam)
{    
    Proxy* proxy = (Proxy*)lpParam;
    return proxy->_workerThreadProc();
}

DWORD Proxy::_workerThreadProc()
{    
    int nBytesRecv = 0;
    int nBytesSent = 0;
    DWORD dwBytes = 0, dwFlags = 0;

    //Worker thread will be around to process requests, until a Shutdown event is not Signaled.
    while( WAIT_OBJECT_0 != WaitForSingleObject(_shutdownEvent, 0) )
    {
        DWORD bytesTransfered = 0;
        OVERLAPPED* overlapped = NULL;
        ULONG_PTR completionKey = NULL;
        BOOL res = GetQueuedCompletionStatus(_iocp, &bytesTransfered, &completionKey, &overlapped, INFINITE);
        /* !!!
        if(completionKey == NULL || overlapped == NULL) {
            //We are shutting down
            break;
        }
        */

        BridgeContext* bridge = (BridgeContext*)overlapped;
        if(res == FALSE || bytesTransfered == 0) {
            _closeBridge(bridge);
            continue;
        }

        int ret;
        switch(bridge->complitedOperation)
        {
        case BridgeContext::IOAccept:
            ret = bridge->getConnectContext().completeAccept(_listenSocket, bytesTransfered);
            if ((SOCKET_ERROR == res) && (WSA_IO_PENDING != WSAGetLastError())) {
                dbgprint("Error occurred while executing bridge.completeAccept(): %d", WSAGetLastError());
                _closeBridge(bridge);
            }
            _createConnection();
            break;

        case BridgeContext::IORead:
            ret = bridge->send(bytesTransfered);
            if ((SOCKET_ERROR == res) && (WSA_IO_PENDING != WSAGetLastError())) {
                dbgprint("Error occurred while executing bridge.send(): %d", WSAGetLastError());
                _closeBridge(bridge);
            }
            break;

        case BridgeContext::IOWrite:
            ret = bridge->recv();
            if ((SOCKET_ERROR == res) && (WSA_IO_PENDING != WSAGetLastError())) {
                dbgprint("Error occurred while executing bridge.recv(): %d", WSAGetLastError());
                _closeBridge(bridge);
            }
            break;
        }
    }

    return 0;
}

ConnectContext* Proxy::_createConnection()
{
    ConnectContext* connect = new ConnectContext();
    if( !connect->initialize(_iocp) || !connect->waitAccept(_listenSocket) ) {
        delete connect;
        return NULL;
    }

    Lock lock(_cs);
    _connections.push_back(connect);

    return connect;
}

void Proxy::_deleteConnection(ConnectContext* connection)
{
    Lock lock(_cs);
    _connections.remove(connection);
    delete connection;
}

void Proxy::_closeBridge(BridgeContext* bridge)
{
    Lock lock(_cs);
    ConnectContext& connect = bridge->getConnectContext();
    if( connect.close(bridge) ) {
        _deleteConnection(&connect);
    }
}

void Proxy::_clearConnections()
{
    Lock lock(_cs);
    for(Connections::iterator it = _connections.begin(); it != _connections.end(); ++it) {
        delete (*it);
    }
    _connections.clear();
}
